package com.itbaizhan.cxpt.service;

import okhttp3.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.JsonNode;

import java.io.File;
import java.io.IOException;
import java.util.*;

public class MoonshotClient {
    private static final String BASE_URL = "https://api.moonshot.cn/v1";
    private static final String API_KEY = "sk-a7XZ21b5j0XJW2jcJTAxth2phMe1MI08EgcooAFvf1Ko8tOm";
    private static final OkHttpClient client = new OkHttpClient();
    private static final ObjectMapper mapper = new ObjectMapper();

    // 上传文件，返回文件内容
    public static List<Map<String, String>> uploadFilesWithPrompt(List<String> filePaths) throws IOException {
        List<Map<String, String>> messages = new ArrayList<>();

        for (String path : filePaths) {
            // 上传文件
            File file = new File(path);
            RequestBody fileBody = RequestBody.create(file, MediaType.parse("application/octet-stream"));

            MultipartBody requestBody = new MultipartBody.Builder()
                    .setType(MultipartBody.FORM)
                    .addFormDataPart("file", file.getName(), fileBody)
                    .addFormDataPart("purpose", "file-extract")
                    .build();

            Request request = new Request.Builder()
                    .url(BASE_URL + "/files")
                    .addHeader("Authorization", "Bearer " + API_KEY)
                    .post(requestBody)
                    .build();

            Response response = client.newCall(request).execute();
            String result = response.body().string();
            response.close();

            JsonNode rootNode = mapper.readTree(result);
            String fileId = rootNode.get("id").asText();

            // 获取文件内容
            Request contentRequest = new Request.Builder()
                    .url(BASE_URL + "/files/" + fileId + "/content")
                    .addHeader("Authorization", "Bearer " + API_KEY)
                    .get()
                    .build();

            Response contentResponse = client.newCall(contentRequest).execute();
            String fileContent = contentResponse.body().string();
            contentResponse.close();

            // 生成带提示词的 message
            Map<String, String> message = Map.of(
                    "role", "system",
                    "content", "以下是一个PDF文件的内容，请根据这些内容提取以下信息：\n" +
                            "1. 标准号（完整的国标编号）\n" +
                            "2. 标准名称（完整的标准名称）\n" +
                            "3. 在PDF中的页码\n" +
                            "请以JSON格式返回，包含standardNo、standardName和page字段。\n\n" +
                            "PDF内容如下：\n" + fileContent
            );

            messages.add(message);
        }

        return messages;
    }


    // 发送对话请求
    public static void sendChatCompletion(List<Map<String, String>> fileMessages) throws IOException {
        // 构造 messages
        List<Map<String, String>> messages = new ArrayList<>(fileMessages);

        // system 提示词
        Map<String, String> systemMessage = Map.of(
                "role", "system",
                "content", "你是一个专门用于分析PDF文件中国标信息的AI助手。请仔细分析内容中的国标信息，提取以下信息：\n" +
                        "1. 标准号（完整的国标编号）\n" +
                        "2. 标准名称（完整的标准名称）\n" +
                        "3. 在PDF中的页码\n" +
                        "请以JSON格式返回，包含standardNo、standardName和page字段。"
        );
        messages.add(systemMessage);

        // 拼接所有文件内容
        StringBuilder contentBuilder = new StringBuilder();
        for (Map<String, String> msg : fileMessages) {
            contentBuilder.append(msg.get("content")).append("\n");
        }

        // user 提示词
        Map<String, String> userMessage = Map.of(
                "role", "user",
                "content", "请分析以下内容中的国标信息：\n" + contentBuilder.toString()
        );
        messages.add(userMessage);

        // 构造请求体
        Map<String, Object> bodyMap = new HashMap<>();
        bodyMap.put("model", "moonshot-v1-128k");
        bodyMap.put("messages", messages);

        String jsonBody = mapper.writeValueAsString(bodyMap);

        RequestBody requestBody = RequestBody.create(
                jsonBody, MediaType.parse("application/json"));

        Request request = new Request.Builder()
                .url(BASE_URL + "/chat/completions")
                .addHeader("Authorization", "Bearer " + API_KEY)
                .post(requestBody)
                .build();

        // 发送请求
        Response response = client.newCall(request).execute();
        String result = response.body().string();
        response.close();

        // 解析返回内容
        JsonNode rootNode = mapper.readTree(result);
        String reply = rootNode.get("choices").get(0).get("message").get("content").asText();

        System.out.println("Kimi 回复：" + reply);
    }

    public static void main(String[] args) throws IOException {
        List<String> filePaths = Arrays.asList("C:\\Users\\Administrator\\Desktop\\92矿物源腐植酸型农用微生物菌剂.pdf");

        List<Map<String, String>> fileMessages = uploadFilesWithPrompt(filePaths);

        sendChatCompletion(fileMessages);
    }
}
